#include <iostream>
#include <iomanip>
#include <sstream>
#include <cstdlib>
#include <map>
#include <algorithm>
#include <fstream>
#include <time.h>

#include "fastjet/ClusterSequence.hh"
#include "fastjet/tools/Pruner.hh"

#include "Njettiness.hh"

#include "CmdLine.hh"
#include "TFile.h"
#include "TChain.h"

#include "TH1F.h"
#include "TH2F.h"

#include "Analysis.h"


class ATLAS_verification : public Analysis
{
public:
  
  vector<PseudoJet> jets_antikt;
  vector<PseudoJet> jets_cambridge;
  vector<PseudoJet> jets_input;

  ATLAS_verification(const CmdLine& cmdline)
    : Analysis(cmdline)
  {
    fileout->cd();
    hist["met"]=new TH1D("MEt", "#slash{E}_{T}", 50, 0, 500);
    hist["nevt"]=new TH1D("nevt", "nevt", 10, 0, 10);
    hist["nevt_unweighted"]=new TH1D("nevt_unweighted", "nevt_unweighted",
				     10, 0, 10);
  }

  void hist_init();
  void analyze();
  void makejetsinput();

};



int main(int argc, char** argv)
{

  //initialize command line interface
  CmdLine cmdline(argc,argv);

  if(!cmdline.present("-i") || !cmdline.present("-o"))
    {
      cout<<"ATLAS_verification: Produce verification plots in 1203.4606";
      cout<<"Usage: -i <input> -o <output> -n [nev] -p -w"<<endl;
      cout<<"-n \t number of events to process"<<endl;
      cout<<"-w \t whether event has a weight stored as evt Number"<<endl;
      
      return 1;
    }

  //create analysis object
  ATLAS_verification myanalysis(cmdline);
  myanalysis.tree_init();
  myanalysis.hist_init();
  myanalysis.analyze();

  return 0;
}

void ATLAS_verification::hist_init()
{
  hist["cambridge_pt"]=new TH1D("cambridge_pt", "CA Jet p_{T} distribution",
				8, 200, 1000);
  
  hist["cambridge_eta"]=new TH1D("cambridge_eta", "CA Jet #eta distribution",
				 10, -2.5, 2.5);
  
  hist["cambridge_mass_bin1"]=new TH1D("cambridge_mass_bin1", 
				       "CA Jet mass distribution 200 < p_{T} < 300",
				       21, 0, 210);

  hist["cambridge_mass_bin2"]=new TH1D("cambridge_mass_bin2", 
				       "CA Jet mass distribution 300 < p_{T} < 400",
				       21, 0, 210);
  
  hist["cambridge_mass_bin3"]=new TH1D("cambridge_mass_bin3", 
				       "CA Jet mass distribution 400 < p_{T} < 500",
				       30, 0, 300);

  hist["cambridge_mass_bin4"]=new TH1D("cambridge_mass_bin4", 
				       "CA Jet mass distribution 500 < p_{T} < 600",
				       30, 0, 300);
  
  hist["cambridge_tau21_bin1"]=new TH1D("cambridge_tau21_bin1", 
				       "CA Jet #tau_{21} distribution 200 < p_{T} < 300",
					15, 0, 1.2);
  
  hist["cambridge_tau21_bin2"]=new TH1D("cambridge_tau21_bin2", 
				       "CA Jet #tau_{21} distribution 300 < p_{T} < 400",
					15, 0, 1.2);

  hist["cambridge_tau21_bin3"]=new TH1D("cambridge_tau21_bin3", 
				       "CA Jet #tau_{21} distribution 400 < p_{T} < 500",
					15, 0, 1.2);

  hist["cambridge_tau21_bin4"]=new TH1D("cambridge_tau21_bin4", 
				       "CA Jet #tau_{21} distribution 500 < p_{T} < 600",
					15, 0, 1.2);


  //tau32
  hist["cambridge_tau32_bin1"]=new TH1D("cambridge_tau32_bin1", 
				       "CA Jet #tau_{32} distribution 200 < p_{T} < 300",
					15, 0, 1.2);

  hist["cambridge_tau32_bin2"]=new TH1D("cambridge_tau32_bin2", 
				       "CA Jet #tau_{32} distribution 300 < p_{T} < 400",
					15, 0, 1.2);

  hist["cambridge_tau32_bin3"]=new TH1D("cambridge_tau32_bin3", 
				       "CA Jet #tau_{32} distribution 400 < p_{T} < 500",
					15, 0, 1.2);

  hist["cambridge_tau32_bin4"]=new TH1D("cambridge_tau32_bin4", 
				       "CA Jet #tau_{32} distribution 500 < p_{T} < 600",
					15, 0, 1.2);




  //antikt jets
  hist["antikt_pt"]=new TH1D("antikt_pt", "antikt Jet p_{T} distribution",
			     8, 200, 1000);
  
  hist["antikt_eta"]=new TH1D("antikt_eta", "antikt Jet #eta distribution",
			      10, -2.5, 2.5);
  
  hist["antikt_mass_bin1"]=new TH1D("antikt_mass_bin1", 
				       "antikt Jet mass distribution 200 < p_{T} < 300",
				    21, 0, 210);

  hist["antikt_mass_bin2"]=new TH1D("antikt_mass_bin2", 
				       "antikt Jet mass distribution 300 < p_{T} < 400",
				    21, 0, 210);
  
  hist["antikt_mass_bin3"]=new TH1D("antikt_mass_bin3", 
				       "antikt Jet mass distribution 400 < p_{T} < 500",
				    30, 0, 300);

  hist["antikt_mass_bin4"]=new TH1D("antikt_mass_bin4", 
				       "antikt Jet mass distribution 500 < p_{T} < 600",
				    30, 0, 300);
  
  hist["antikt_tau21_bin1"]=new TH1D("antikt_tau21_bin1", 
				       "antikt Jet #tau_{21} distribution 200 < p_{T} < 300",
				     15, 0, 1.2);
  
  hist["antikt_tau21_bin2"]=new TH1D("antikt_tau21_bin2", 
				       "antikt Jet #tau_{21} distribution 300 < p_{T} < 400",
				     15, 0, 1.2);

  hist["antikt_tau21_bin3"]=new TH1D("antikt_tau21_bin3", 
				       "antikt Jet #tau_{21} distribution 400 < p_{T} < 500",
				     15, 0, 1.2);

  hist["antikt_tau21_bin4"]=new TH1D("antikt_tau21_bin4", 
				       "antikt Jet #tau_{21} distribution 500 < p_{T} < 600",
				     15, 0, 1.2);


  //tau32
  hist["antikt_tau32_bin1"]=new TH1D("antikt_tau32_bin1", 
				       "antikt Jet #tau_{32} distribution 200 < p_{T} < 300",
				     15, 0, 1.2);

  hist["antikt_tau32_bin2"]=new TH1D("antikt_tau32_bin2", 
				       "antikt Jet #tau_{32} distribution 300 < p_{T} < 400",
				     15, 0, 1.2);

  hist["antikt_tau32_bin3"]=new TH1D("antikt_tau32_bin3", 
				       "antikt Jet #tau_{32} distribution 400 < p_{T} < 500",
				     15, 0, 1.2);

  hist["antikt_tau32_bin4"]=new TH1D("antikt_tau32_bin4", 
				       "antikt Jet #tau_{32} distribution 500 < p_{T} < 600",
				     15, 0, 1.2);

}

void ATLAS_verification::makejetsinput()
{
  jets_input.clear();
  jets_antikt.clear();
  jets_cambridge.clear();
  
  for(Int_t i=0; i<size; i++)
    if(Pt[i]>1e-20)
      {
	if(Status[i]!=1)
	  continue;
	
	if(abs(PID[i])==12 || 
	   abs(PID[i])==14 || 
	   abs(PID[i])==16)
	  continue;
	
	PseudoJet j(Px[i], Py[i], 
		    Pz[i], E[i]);
	
	if(fabs(Eta[i])>3)
	  continue;
	
	  if(fabs(Eta[i])>2.5)
	    if(PID[i]==22 ||
	       abs(PID[i])==11 || 
	       abs(PID[i])==13)
	      continue;
	  
	  jets_input.push_back(j);
      }
}



//main function to loop over all the entries
void ATLAS_verification::analyze()
{
  //begin event loop
  for(int entry=0; 
      (maxevt<0 || entry<maxevt) && entry<tree.GetEntries();
      entry++)
    {

      if(entry%1000==0 && entry/1000 > 0)
	cout<<"MESSAGE: "<<entry<<" Processed"<<endl;

      //populate our variables
      tree.GetEntry(entry);

      Event_=entry;
      if(cmdline.present("-unit-weight"))
	{weight=1.0;}

      if(cmdline.present("-weightcut"))
	{
	  if(entry==0)
	    cout<<"MESSAGE: weightcut enabled "<<cmdline.value<double>("-weightcut")<<endl;
	  if (weight > cmdline.value<double>("-weightcut")*3.35489e7)
	    continue;
	}

      hist["nevt"]->Fill(0.0, weight);
      hist["nevt_unweighted"]->Fill(0.0);
      
      

      //grab jets constiuents
      makejetsinput();

      using namespace fastjet;
      //get cambridge jets and antikt jets
      JetDefinition jetdef_cambridge(cambridge_algorithm, 1.2);
      JetDefinition jetdef_antikt(antikt_algorithm, 1.0);
      
      ClusterSequence* clust_cambridge
	=new ClusterSequence(jets_input, jetdef_cambridge);
      
      ClusterSequence* clust_antikt
	=new ClusterSequence(jets_input, jetdef_antikt);
      
      //only care about jets with minimum 200 pT
      Selector select_pt_eta 
	= SelectorAbsRapMax(2.0) && SelectorPtMin(200.0);
      
      jets_cambridge
	= sorted_by_pt(select_pt_eta(clust_cambridge->inclusive_jets()));

      jets_antikt
	= sorted_by_pt(select_pt_eta(clust_antikt->inclusive_jets()));

      //initialize Nsubjettiness
      //beta, R_0, normalization
      NsubParameters param_cambridge(1.0, 1.2, 50.);
      NsubParameters param_antikt(1.0, 1.0, 50.);
      

      Njettiness njettiness_cambridge(Njettiness::kt_axes,
				      param_cambridge);

      Njettiness njettiness_antikt(Njettiness::kt_axes,
				   param_antikt);
      

      for(int i=0; i<jets_cambridge.size(); i++)
	{
	  double c_pt=jets_cambridge[i].perp();
	  hist["cambridge_pt"]->Fill(c_pt, weight);
	  hist["cambridge_eta"]->Fill(jets_cambridge[i].eta(), weight);
	  
	  if(c_pt < 200. || c_pt > 600.) continue;

	  //plot subjettiness
	  vector<PseudoJet> constituents
	    = jets_cambridge[i].constituents();
	  double tau1= njettiness_cambridge.getTau(1, constituents);
	  double tau2= njettiness_cambridge.getTau(2, constituents);
	  double tau3= njettiness_cambridge.getTau(3, constituents);
	  
	  //bin1
	  if(c_pt < 300.)
	    {
	      hist["cambridge_mass_bin1"]->Fill(jets_cambridge[i].m(), weight);
	      hist["cambridge_tau21_bin1"]->Fill(tau2/tau1, weight);
	      hist["cambridge_tau32_bin1"]->Fill(tau3/tau2, weight);
	    }
	  
	  else if(c_pt < 400.)
	    {
	      hist["cambridge_mass_bin2"]->Fill(jets_cambridge[i].m(), weight);
	      hist["cambridge_tau21_bin2"]->Fill(tau2/tau1, weight);
	      hist["cambridge_tau32_bin2"]->Fill(tau3/tau2, weight);
	    }

	  else if(c_pt < 500.)
	    {
	      hist["cambridge_mass_bin3"]->Fill(jets_cambridge[i].m(), weight);
	      hist["cambridge_tau21_bin3"]->Fill(tau2/tau1, weight);
	      hist["cambridge_tau32_bin3"]->Fill(tau3/tau2, weight);
	    }

	  else
	    {
	      hist["cambridge_mass_bin4"]->Fill(jets_cambridge[i].m(), weight);
	      hist["cambridge_tau21_bin4"]->Fill(tau2/tau1, weight);
	      hist["cambridge_tau32_bin4"]->Fill(tau3/tau2, weight);
	    }
	}


      for(int i=0; i<jets_antikt.size(); i++)
	{
	  double c_pt=jets_antikt[i].perp();
	  hist["antikt_pt"]->Fill(c_pt, weight);
	  hist["antikt_eta"]->Fill(jets_antikt[i].eta(), weight);
	  
	  if(c_pt < 200. || c_pt > 600.) continue;

	  //plot subjettiness
	  vector<PseudoJet> constituents
	    = jets_antikt[i].constituents();
	  double tau1= njettiness_antikt.getTau(1, constituents);
	  double tau2= njettiness_antikt.getTau(2, constituents);
	  double tau3= njettiness_antikt.getTau(3, constituents);
	  
	  //bin1
	  if(c_pt < 300.)
	    {
	      hist["antikt_mass_bin1"]->Fill(jets_antikt[i].m(), weight);
	      hist["antikt_tau21_bin1"]->Fill(tau2/tau1, weight);
	      hist["antikt_tau32_bin1"]->Fill(tau3/tau2, weight);
	    }
	  
	  else if(c_pt < 400.)
	    {
	      hist["antikt_mass_bin2"]->Fill(jets_antikt[i].m(), weight);
	      hist["antikt_tau21_bin2"]->Fill(tau2/tau1, weight);
	      hist["antikt_tau32_bin2"]->Fill(tau3/tau2, weight);
	    }

	  else if(c_pt < 500.)
	    {
	      hist["antikt_mass_bin3"]->Fill(jets_antikt[i].m(), weight);
	      hist["antikt_tau21_bin3"]->Fill(tau2/tau1, weight);
	      hist["antikt_tau32_bin3"]->Fill(tau3/tau2, weight);
	    }

	  else
	    {
	      hist["antikt_mass_bin4"]->Fill(jets_antikt[i].m(), weight);
	      hist["antikt_tau21_bin4"]->Fill(tau2/tau1, weight);
	      hist["antikt_tau32_bin4"]->Fill(tau3/tau2, weight);
	    }
	}      

      delete clust_antikt;
      delete clust_cambridge;
    }
  
  fileout->Write();
  fileout->Close();
}


